#include "stdafx.h"
#include "DriveEvent.h"
#include "global.h"

//
//	API events handler
//

bool DriveEventHandler::init()
{
    bool result = false;

    for (;;)
    {
#ifdef _CFG_UDP_
        if (!m_udpProxy.init(this,
                             (char *)g_proxyAddress,
                             (((sockaddr *)g_proxyAddress)->sa_family == AF_INET6) ? sizeof(sockaddr_in6) : sizeof(sockaddr_in),
                             g_userName.empty() ? NULL : g_userName.c_str(),
                             g_userPassword.empty() ? NULL : g_userPassword.c_str()))
        {
            DbgPrint("Unable to start UDP proxy");
            break;
        }
#endif
        if (!m_tcpProxy.init(htons(8888)))
        {
            DbgPrint("Unable to start TCP proxy");
            break;
        }

        m_tcpProxy.setProxy(0, TcpProxy::PROXY_SOCKS5,
                            (char *)g_proxyAddress,
                            (((sockaddr *)g_proxyAddress)->sa_family == AF_INET6) ? sizeof(sockaddr_in6) : sizeof(sockaddr_in),
                            g_userName.empty() ? NULL : g_userName.c_str(),
                            g_userPassword.empty() ? NULL : g_userPassword.c_str());

        result = true;
        m_tcpProxy.setEventHandler(&m_proxyHandler);
        m_proxyHandler.senEventhandle(&m_tcpProxy, &gUserMgr);
        break;
    }

    if (!result)
    {
        free();
    }

    return result;
}

void DriveEventHandler::free()
{
#ifdef _CFG_UDP_
    m_udpProxy.free();
#endif
    m_tcpProxy.free();

    AutoLock lock(m_cs);
    while (!m_udpCtxMap.empty())
    {
        tUdpCtxMap::iterator it = m_udpCtxMap.begin();
        delete it->second;
        m_udpCtxMap.erase(it);
    }
    m_filteredUdpIds.clear();
}
void DriveEventHandler::onUdpReceiveComplete(ENDPOINT_ID id,
                                             char *buf, int len, char *remoteAddress,
                                             int remoteAddressLen)
{
#ifdef _CFG_UDP_
    AutoLock lock(m_cs);

    tUdpCtxMap::iterator it = m_udpCtxMap.find(id);
    if (it == m_udpCtxMap.end())
        return;

    nf_udpPostReceive(id, (const unsigned char *)remoteAddress, buf, len, it->second->m_options);
#endif

#ifdef UDP_DEBUG
    char remoteAddr[MAX_PATH];
    DWORD dwLen;

    dwLen = sizeof(remoteAddr);
    WSAAddressToString((sockaddr *)remoteAddress,
                       (((sockaddr *)remoteAddress)->sa_family == AF_INET6) ? sizeof(sockaddr_in6) : sizeof(sockaddr_in),
                       NULL,
                       remoteAddr,
                       &dwLen);

    DbgPrint("onUdpReceiveComplete id=%I64u len=%d remoteAddress=%s\n", id, len, remoteAddr);
#endif // _DEBUG

}

void DriveEventHandler::threadStart()
{
}

void DriveEventHandler::threadEnd()
{
}

//
// TCP events
//
void DriveEventHandler::tcpConnectRequest(ENDPOINT_ID id, PNF_TCP_CONN_INFO pConnInfo)
{
    DbgPrint("tcpConnectRequest id=%I64u", id);

    sockaddr *pAddr = (sockaddr *)pConnInfo->remoteAddress;
    int addrLen = (pAddr->sa_family == AF_INET6) ? sizeof(sockaddr_in6) : sizeof(sockaddr_in);

    // Don't redirect the connection if it is already redirected
    if (memcmp(pAddr, g_proxyAddress, addrLen) == 0)
    {
        DbgPrint("tcpConnectRequest id=%I64u bypass already redirected", id);
        return;
    }

    if (g_processNames.size() > 0)
    {
        if (!checkProcessName(pConnInfo->processId))
        {
            DbgPrint("tcpConnectRequest id=%I64u bypass wrong process", id);
            return;
        }
    }

    if (!m_tcpProxy.isIPFamilyAvailable(pAddr->sa_family))
    {
        DbgPrint("tcpConnectRequest id=%I64u bypass ipFamily %d", id, pAddr->sa_family);
        return;
    }

    m_tcpProxy.setConnInfo(pConnInfo);

    // Redirect the connection
    if (pAddr->sa_family == AF_INET)
    {
        sockaddr_in addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin_family = AF_INET;
        addr.sin_addr.S_un.S_addr = inet_addr("127.0.0.1");
        addr.sin_port = m_tcpProxy.getPort();

        memcpy(pConnInfo->remoteAddress, &addr, sizeof(addr));
    }
    else
    {
        sockaddr_in6 addr;
        memset(&addr, 0, sizeof(addr));
        addr.sin6_family = AF_INET6;
        addr.sin6_addr.u.Byte[15] = 1;
        addr.sin6_port = m_tcpProxy.getPort();

        memcpy(pConnInfo->remoteAddress, &addr, sizeof(addr));
    }

    // Specify current process id to avoid blocking connection redirected to local proxy
    pConnInfo->processId = GetCurrentProcessId();
}

void DriveEventHandler::tcpConnected(ENDPOINT_ID id, PNF_TCP_CONN_INFO pConnInfo)
{
    DbgPrintConnInfo(true, id, pConnInfo);
}

void DriveEventHandler::tcpClosed(ENDPOINT_ID id, PNF_TCP_CONN_INFO pConnInfo)
{
    DbgPrintConnInfo(false, id, pConnInfo);
}

void DriveEventHandler::tcpReceive(ENDPOINT_ID id, const char *buf, int len)
{
    DbgPrint("tcpReceive id=%I64u len=%d", id, len);
    // Send the packet to application
    nf_tcpPostReceive(id, buf, len);
}

void DriveEventHandler::tcpSend(ENDPOINT_ID id, const char *buf, int len)
{
    DbgPrint("tcpSend id=%I64u len=%d", id, len);
    // Send the packet to server
    nf_tcpPostSend(id, buf, len);
}

void DriveEventHandler::tcpCanReceive(ENDPOINT_ID id)
{
}

void DriveEventHandler::tcpCanSend(ENDPOINT_ID id)
{
}

//
// UDP events
//

void DriveEventHandler::udpCreated(ENDPOINT_ID id, PNF_UDP_CONN_INFO pConnInfo)
{
#ifdef _CFG_UDP_
    DbgPrintAddrInfo(true, id, pConnInfo);
    if (g_processNames.size() > 0)
    {
        if (checkProcessName(pConnInfo->processId))
        {
            AutoLock lock(m_cs);
            m_filteredUdpIds.insert(id);
        }
    }
    else
    {
        AutoLock lock(m_cs);
        m_filteredUdpIds.insert(id);
    }
#endif // _CFG_UDP_
}

void DriveEventHandler::udpConnectRequest(ENDPOINT_ID id, PNF_UDP_CONN_REQUEST pConnReq)
{
#ifdef UDP_DEBUG
    DbgPrint("udpConnectRequest id=%I64u\n", id);
#endif // UDP_DEBUG
}

void DriveEventHandler::udpClosed(ENDPOINT_ID id, PNF_UDP_CONN_INFO pConnInfo)
{
#ifdef UDP_DEBUG
    DbgPrintAddrInfo(false, id, pConnInfo);
#endif

#ifdef _CFG_UDP_
    m_udpProxy.deleteProxyConnection(id);

    AutoLock lock(m_cs);

    tUdpCtxMap::iterator it = m_udpCtxMap.find(id);
    if (it != m_udpCtxMap.end())
    {
        delete it->second;
        m_udpCtxMap.erase(it);
    }

    m_filteredUdpIds.erase(id);
#endif
}

void DriveEventHandler::udpReceive(ENDPOINT_ID id, const unsigned char *remoteAddress,
                                   const char *buf, int len, PNF_UDP_OPTIONS options)
{
#ifdef UDP_DEBUG
    char remoteAddr[MAX_PATH];
    DWORD dwLen;

    dwLen = sizeof(remoteAddr);
    WSAAddressToString((sockaddr *)remoteAddress,
                       (((sockaddr *)remoteAddress)->sa_family == AF_INET6) ? sizeof(sockaddr_in6) : sizeof(sockaddr_in),
                       NULL,
                       remoteAddr,
                       &dwLen);

    DbgPrint("udpReceive id=%I64u len=%d remoteAddress=%s\n", id, len, remoteAddr);
#endif
    // Send the packet to application
    nf_udpPostReceive(id, remoteAddress, buf, len, options);
}

void DriveEventHandler::udpSend(ENDPOINT_ID id, const unsigned char *remoteAddress,
                                const char *buf, int len, PNF_UDP_OPTIONS options)
{
#ifdef UDP_DEBUG
    char remoteAddr[MAX_PATH];
    DWORD dwLen;

    dwLen = sizeof(remoteAddr);
    WSAAddressToString((sockaddr *)remoteAddress,
                       (((sockaddr *)remoteAddress)->sa_family == AF_INET6) ? sizeof(sockaddr_in6) : sizeof(sockaddr_in),
                       NULL,
                       remoteAddr,
                       &dwLen);

    DbgPrint("udpSend id=%I64u len=%d remoteAddress=%s", id, len, remoteAddr);
#endif // _DEBUG

#ifdef _CFG_UDP_
    {
        AutoLock lock(m_cs);

        tIdSet::iterator itid = m_filteredUdpIds.find(id);
        if (itid == m_filteredUdpIds.end())
        {
            nf_udpPostSend(id, remoteAddress, buf, len, options);
            return;
        }

        tUdpCtxMap::iterator it = m_udpCtxMap.find(id);
        if (it == m_udpCtxMap.end())
        {
            if (!m_udpProxy.createProxyConnection(id))
                return;

            m_udpCtxMap[id] = new UDP_CONTEXT(options);
        }
    }

    {
        int addrLen = (((sockaddr *)remoteAddress)->sa_family == AF_INET) ? sizeof(sockaddr_in) : sizeof(sockaddr_in6);
        if (!m_udpProxy.udpSend(id, (char *)buf, len, (char *)remoteAddress, addrLen))
        {
            nf_udpPostSend(id, remoteAddress, buf, len, options);
        }
    }
#else
    nf_udpPostSend(id, remoteAddress, buf, len, options);
#endif
}

void DriveEventHandler::udpCanReceive(ENDPOINT_ID id)
{
}

void DriveEventHandler::udpCanSend(ENDPOINT_ID id)
{
}